package com.wenbing.miio.mi;

import com.wenbing.miio.core.*;
import com.wenbing.miio.core.exception.DataFormatException;
import com.wenbing.miio.core.exception.NetworkLayerException;
import com.wenbing.miio.util.BytesUtils;
import com.wenbing.miio.util.CipherUtil;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;

/**
 * 数据帧层，封装/解包 miio协议, 不涉及加解密
 * @author wen bing
 * @projectName: miio
 * @package: com.wenbing.miio.core
 * @className: A
 * @date 2024/09/19 14:08
 **/
public class MiFrameLayer extends AbstractNetworkLayer {

    public static final String ATTR_TIMESTAMP = "timestamp";

    public static final String ATTR_DID = "did";

    public static final String ATTR_FLAGS = "flags";

    public static final String ATTR_TOKEN = "token";

    @Override
    public ByteBuffer unpacking(NetworkLayerContext context, ByteBuffer packingData) throws IOException, NetworkLayerException {
        // 校验报文头
        checkHeader(packingData);
        MessageFrame frame = parseFrame(packingData);
        LOGGER.debug("unpacking flags is {}; did is {}; timestamp is {}", frame.getFlags(), frame.getDeviceId(), frame.getStamp());
        // 数据发送时间
        context.setAttribute(ATTR_TIMESTAMP, frame.getStamp());
        context.setAttribute(ATTR_DID, frame.getDeviceId());
        context.setAttribute(ATTR_FLAGS, frame.getFlags());
        byte[] token = MiDeviceUtil.getToken(this);
        // 校验数据
        checkData(packingData, frame, token);
        return frame.getPayload();
    }

    @Override
    public ByteBuffer packing(NetworkLayerContext context, ByteBuffer unpackingData) throws IOException, NetworkLayerException {
        int did = context.getAttribute(ATTR_DID);
        int flags = context.getAttribute(ATTR_FLAGS, 0);
        long timestamp = context.getAttribute(ATTR_TIMESTAMP);
        byte[] token = MiDeviceUtil.getToken(this);
        LOGGER.debug("packing flags is {}; did is {}; timestamp is {}", flags, did, timestamp);
        return toBuf(unpackingData, flags, did, timestamp, token).flip();
    }

    private ByteBuffer toBuf(ByteBuffer unpackingData, int flags, int did, long timestamp, byte[] token) throws NetworkLayerException {
        ByteBuffer buffer = ByteBuffer.allocate(unpackingData.limit() + MessageFrame.HEAD_LENGTH);
        buffer.putShort(MessageFrame.MAGIC);
        buffer.putShort((short) buffer.capacity());
        buffer.putInt(flags);
        buffer.putInt(did);

        buffer.putInt((int) timestamp);
        // 记录当前位置
        int pos = buffer.position();
        buffer.put(token);
        buffer.put(unpackingData);
        // 重置位置，计算checksum
        buffer.position(0);
        byte[] checksum = unpackingData.limit() == 0 ?
                token : checksum(buffer);
        // 恢复位置，塞入checksum
        buffer.position(pos);
        buffer.put(checksum);
        // 移动到最后
        buffer.position(buffer.capacity());
        return buffer;
    }

    private byte[] checksum(ByteBuffer byteBuffer) throws NetworkLayerException {
        try {
            int offset = byteBuffer.arrayOffset();
            byte[] data = Arrays.copyOfRange(byteBuffer.array(), offset + byteBuffer.position(), offset + byteBuffer.limit());
            return CipherUtil.md5(data);
        } catch (NoSuchAlgorithmException e) {
            throw new NetworkLayerException(e);
        }
    }

    private void checkHeader(ByteBuffer byteBuffer) throws DataFormatException {
        short magic = byteBuffer.getShort();
        if (magic != MessageFrame.MAGIC ){
            // 魔数不匹配
            throw new DataFormatException(String.format("magic[%d] do not match, expect %d", magic, MessageFrame.MAGIC));
        }
        byteBuffer.mark();
        int length = Short.toUnsignedInt(byteBuffer.getShort());
        if (length < MessageFrame.HEAD_LENGTH){
            // 数据错误
            throw new DataFormatException("bad data");
        }
        byteBuffer.reset();
    }

    private void checkData(ByteBuffer byteBuffer, MessageFrame frame, byte[] token) throws DataFormatException, NetworkLayerException {
        byteBuffer.position(MessageFrame.HEAD_LENGTH - MessageFrame.CHECKSUM_LENGTH);
        byteBuffer.put(token);
        byteBuffer.position(0);
        byte[] checksum = checksum(byteBuffer);
        if (Arrays.equals(frame.getChecksum() , checksum)){
            return;
        }
        if (frame.getPayload().limit() == 0 && (
                Arrays.equals(BytesUtils.repeat((byte) 0xff, MessageFrame.CHECKSUM_LENGTH) , frame.getChecksum())
                || Arrays.equals(token, frame.getChecksum()))) {
            return;
        }
        if (frame.getFlags() == MessageFrame.FLAGS_HELLO) {
            return;
        }
        // 数据可能被篡改
        throw new DataFormatException("data may be tampered with");
    }

    public MessageFrame parseFrame(ByteBuffer byteBuffer) throws DataFormatException {
        MessageFrame frame = new MessageFrame();
        int length = Short.toUnsignedInt(byteBuffer.getShort());
        int flags = byteBuffer.getInt();
        frame.setFlags(flags);

        int deviceId = byteBuffer.getInt();
        frame.setDeviceId(deviceId);

        int stamp = byteBuffer.getInt();
        frame.setStamp(Integer.toUnsignedLong(stamp));

        byte[] checkSum = new byte[MessageFrame.CHECKSUM_LENGTH];
        byteBuffer.get(checkSum);
        frame.setChecksum(checkSum);

        ByteBuffer payload = byteBuffer.slice();
        payload.limit(length - MessageFrame.HEAD_LENGTH);
        frame.setPayload(payload);
        byteBuffer.limit(length);
        return  frame;
    }
}
